# coding=utf-8
"""
    token 机制
"""

import jwt
import pytz
import datetime
from flask import request, current_app
from utils.token_data_util import token_store

# @Time    :  2024-10-16 17:31:38
# @Author  :  jiangtong
# @Email   :  gingerqgyy@outlook.com
# @Project :  ginger_yl_test_boost
# @File    :  token_utils


token_store = token_store
# 设置北京时间的时区
BEIJING_TZ = pytz.timezone("Asia/Shanghai")

token_res = {

    # Js 跳转登录 : ( 访问 token 无效, 刷新 token 过期, 访问 token 缺失, 访问 token 注销, 刷新 token 缺失 )
    'access_failure': ('访问 token 无效', 'ACCESS_TOKEN_FAILURE'),  # 访问 token 无效, ( 假的, 错的 )
    'refresh_expire': ('刷新 token 过期', 'REFRESH_TOKEN_FAILURE'),  # 刷新 token 过期 ( 过期的 )
    'access_lose': ('访问 token 缺失', 'ACCESS_TOKEN_LOSE'),  # 访问 token 缺失 ( 缺失, 未携带 )
    'access_removed': ('访问 token 注销', 'ACCESS_TOKEN_REMOVED'),  # 访问 token 清除 ( 用户注销登录, 内存中干掉了 )
    'refresh_lose': ('刷新 token 缺失', 'REFRESH_TOKEN_LOSE'),  # 刷新 token 缺失
    'refresh_failure': ('刷新 token 无效', 'REFRESH_TOKEN_FAILURE'),  # 刷新 token 无效 ( 当 JS 根据刷新 token 获取访问 token 时, 刷新 token 是一个无效的 token )

    # Js 重新获取访问 token :
    'refresh': ('访问 token 失效', 'ACCESS_TOKEN_EXPIRE'),  # 访问 token 失效, ( 根据刷新 token 生成新的访问 token )

    # 更新成功 :
    'update_access': ('更新 access_token 成功', 'ACCESS_TOKEN_UPDATED')
}


# 生成 token 函数 ( 访问 token 和刷新 token ) :
def generate_token(username, update_refresh=True):
    """ 生成初始访问 token 和刷新 token """

    # 获取 token_version :
    token_version = token_store.get(username).get('token_version') + 1 if token_store.get(username) else 1

    # 访问 token 过期时间 ( 例如 2 小时 ) :
    expiration = datetime.datetime.now(BEIJING_TZ) + datetime.timedelta(hours=current_app.config['JWT_EXPIRATION_HOURS'])
    # 刷新 token 过期时间 ( 例如 7 天 ) :
    refresh_expiration = datetime.datetime.now(BEIJING_TZ) + datetime.timedelta(days=current_app.config['JWT_REFRESH_DAYS'])

    # 生成访问 token
    access_token = jwt.encode(
        {'username': username, 'exp': expiration.timestamp(), 'token_version': token_version},
        current_app.config['SECRET_KEY'], algorithm=current_app.config['JWT_ALGORITHM']
    )

    # 生成刷新 token
    refresh_token = token_store.get(username).get('refresh_token') if token_store.get(username) else None
    if update_refresh:
        refresh_token = jwt.encode(
            {'username': username, 'exp': refresh_expiration.timestamp(), 'token_version': token_version},
            current_app.config['SECRET_KEY'], algorithm=current_app.config['JWT_ALGORITHM']
        )

    # 存储访问 token 和刷新 token 到内存 :
    token_store[username] = {
        'token_version': token_version,
        'access_token': access_token,
        'refresh_token': refresh_token,
        'expiration': expiration.timestamp(),
        'refresh_expiration': refresh_expiration.timestamp()
    }

    # 更新 refresh_token : ( 当通过 refresh_token 更新 access_token 时, 不更新 refresh_token )
    if update_refresh:
        token_store[username]['refresh_token'] = refresh_token

    return access_token, refresh_token


# 校验 token 并刷新
def verify_token(token):
    """ 验证访问 token 并自动刷新过期时间 """

    try:
        # 解析访问 token : 能提取到 username, exp ( 过期时间 )
        decoded = jwt.decode(token, current_app.config['SECRET_KEY'], algorithms=[current_app.config['JWT_ALGORITHM']])
        now = datetime.datetime.now(BEIJING_TZ)

        username = decoded['username']
        token_data = token_store.get(username)
        if not token_data:
            return False, None, token_res.get('access_removed')

        in_version = decoded['token_version']
        on_version = token_data.get('token_version')
        if in_version != on_version:
            return False, None, token_res.get('access_failure')

        # 检查访问 token 是否接近过期 ( 例如剩 1 小时 ) :
        exp_datetime = datetime.datetime.fromtimestamp(decoded['exp'], tz=BEIJING_TZ)
        if exp_datetime < now + datetime.timedelta(hours=1):

            # 获取刷新 token
            refresh_token = token_data['refresh_token']

            try:
                # 解析刷新 token : ( 解析刷新 token 也会返回数据, 也是包含 username 和 exp, 刷新 token 过期后, 解析会抛异常 )
                jwt.decode(refresh_token, current_app.config['SECRET_KEY'], algorithms=[current_app.config['JWT_ALGORITHM']])

                # 刷新 token 仍然有效, 生成新的访问 token ( 放行, 响应头中携带新 token )
                new_access_token, _ = generate_token(username)
                # return True, username, new_access_token
                return True, new_access_token, None

            # 刷新 token 已过期
            except jwt.ExpiredSignatureError:
                return False, None, token_res.get('refresh_expire')
            # 刷新 token 无效
            except jwt.InvalidTokenError:
                return False, None, token_res.get('refresh_failure')

        # 访问 token 有效且剩余时间大于 1 小时 : ( 放行 )
        return True, None, None

    # 访问 token 过期 : ( 需要调用刷新 token 进行刷新 )
    except jwt.ExpiredSignatureError:
        return True, None, token_res.get('refresh')
    # 访问 token 无效 : ( 跳转登录页, 引导登录 )
    except jwt.InvalidTokenError:
        return False, None, token_res.get('access_failure')


# 通过 refresh_token 刷新 access_token:
def refresh_access_token():
    try:
        token = request.headers.get('Refresh-Token')
        if not token:
            return False, token_res.get('refresh_lose')
        refresh_decode = jwt.decode(token, current_app.config['SECRET_KEY'], algorithms=[current_app.config['JWT_ALGORITHM']])
        # username, _ = jwt.decode(token, current_app.config['SECRET_KEY'], algorithms=[current_app.config['JWT_ALGORITHM']])
        username = refresh_decode.get('username')
        access_token, __ = generate_token(username=username, update_refresh=False)
        return True, access_token,
    except jwt.ExpiredSignatureError:
        return False, token_res.get('refresh_expire')
    except jwt.InvalidTokenError:
        return False, token_res.get('refresh_failure')


# token 校验与生成新 token
def token_required():
    """ 检查请求中的 token 并处理过期逻辑 """
    current_app.logger.info(f'Current Request Endpoint: {request.endpoint}')

    # 允许无需验证的路由
    if 'login' in request.endpoint or 'static' in request.endpoint or 'refresh_token' in request.endpoint:
        return True, None, None

    # 获取请求头中的 token
    token = request.headers.get('Access-Token')
    if not token:
        return False, None, token_res.get('access_lose')

    """
    username, new_token = verify_token(token)
    if not username:
        return False, None, res_un_auth()

    return True, new_token, None
    """
    # 校验 token 并处理刷新
    return verify_token(token)


if __name__ == '__main__':
    import app

    _app = app.create_app('dev')
    with _app.app_context():
        print(generate_token('ginger'))

        local_token = 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VybmFtZSI6ImdpbmdlciIsImV4cCI6MTcyOTIyNzk3NC4zODQ1NDN9.nQF9-0wlqp8XdBBKdkMkBR-GAfegsbmTWHWNAzKAVfo'
        print(verify_token(local_token))
